/*
 * Copyright (C) 2016 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.androidx.media3.exoplayer.extractor.ts;

import static androidx.media3.extractor.ts.TsPayloadReader.EsInfo.AUDIO_TYPE_UNDEFINED;
import static androidx.media3.extractor.ts.TsPayloadReader.FLAG_PAYLOAD_UNIT_START_INDICATOR;

import android.util.SparseArray;
import android.util.SparseBooleanArray;
import android.util.SparseIntArray;

import androidx.annotation.IntDef;
import androidx.annotation.Nullable;
import androidx.media3.common.C;
import androidx.media3.common.ParserException;
import androidx.media3.common.util.Assertions;
import androidx.media3.common.util.ParsableBitArray;
import androidx.media3.common.util.ParsableByteArray;
import androidx.media3.common.util.TimestampAdjuster;
import androidx.media3.common.util.UnstableApi;
import androidx.media3.common.util.Util;
import androidx.media3.extractor.Extractor;
import androidx.media3.extractor.ExtractorInput;
import androidx.media3.extractor.ExtractorOutput;
import androidx.media3.extractor.ExtractorsFactory;
import androidx.media3.extractor.PositionHolder;
import androidx.media3.extractor.SeekMap;
import androidx.media3.extractor.TrackOutput;
import androidx.media3.extractor.ts.DefaultTsPayloadReaderFactory;
import androidx.media3.extractor.ts.SectionPayloadReader;
import androidx.media3.extractor.ts.SectionReader;
import androidx.media3.extractor.ts.TsPayloadReader;
import androidx.media3.extractor.ts.TsUtil;

import java.io.IOException;
import java.lang.annotation.Documented;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

@SuppressWarnings("ALL")
@UnstableApi
/** Extracts data from the MPEG-2 TS container format. */
public final class MyTsExtractor implements Extractor {

    /**
     * Factory for {@link MyTsExtractor} instances.
     */
    public static final ExtractorsFactory FACTORY = () -> new Extractor[]{new MyTsExtractor()};
    /**
     * Behave as defined in ISO/IEC 13818-1.
     */
    public static final int MODE_MULTI_PMT = 0;
    /**
     * Assume only one PMT will be contained in the stream, even if more are declared by the PAT.
     */
    public static final int MODE_SINGLE_PMT = 1;
    /**
     * Enable single PMT mode, map {@link TrackOutput}s by their type (instead of PID) and ignore
     * continuity counters.
     */
    public static final int MODE_HLS = 2;
    public static final int TS_PACKET_SIZE = 188;
    public static final int DEFAULT_TIMESTAMP_SEARCH_BYTES = 600 * TS_PACKET_SIZE;
    public static final int TS_STREAM_TYPE_MPA = 0x03;
    public static final int TS_STREAM_TYPE_MPA_LSF = 0x04;
    public static final int TS_STREAM_TYPE_AAC_ADTS = 0x0F;
    public static final int TS_STREAM_TYPE_AAC_LATM = 0x11;
    public static final int TS_STREAM_TYPE_AC3 = 0x81;
    public static final int TS_STREAM_TYPE_DTS = 0x8A;
    public static final int TS_STREAM_TYPE_HDMV_DTS = 0x82;
    public static final int TS_STREAM_TYPE_E_AC3 = 0x87;
    public static final int TS_STREAM_TYPE_AC4 = 0xAC; // DVB/ATSC AC-4 Descriptor
    public static final int TS_STREAM_TYPE_H262 = 0x02;
    public static final int TS_STREAM_TYPE_H263 = 0x10; // MPEG-4 Part 2 and H.263
    public static final int TS_STREAM_TYPE_H264 = 0x1B;
    public static final int TS_STREAM_TYPE_H265 = 0x24;
    public static final int TS_STREAM_TYPE_ID3 = 0x15;
    public static final int TS_STREAM_TYPE_SPLICE_INFO = 0x86;
    public static final int TS_STREAM_TYPE_DVBSUBS = 0x59;
    // Stream types that aren't defined by the MPEG-2 TS specification.
    public static final int TS_STREAM_TYPE_DC2_H262 = 0x80;
    public static final int TS_STREAM_TYPE_AIT = 0x101;
    public static final int TS_SYNC_BYTE = 0x47; // First byte of each TS packet.
    private static final int TS_PAT_PID = 0;
    private static final int MAX_PID_PLUS_ONE = 0x2000;
    private static final long AC3_FORMAT_IDENTIFIER = 0x41432d33;
    private static final long E_AC3_FORMAT_IDENTIFIER = 0x45414333;
    private static final long AC4_FORMAT_IDENTIFIER = 0x41432d34;
    private static final long HEVC_FORMAT_IDENTIFIER = 0x48455643;
    private static final int BUFFER_SIZE = TS_PACKET_SIZE * 50;
    private static final int SNIFF_TS_PACKET_COUNT = 5;
    @Mode
    private final int mode;
    private final int timestampSearchBytes;
    private final List<TimestampAdjuster> timestampAdjusters;
    private final ParsableByteArray tsPacketBuffer;
    private final SparseIntArray continuityCounters;
    private final TsPayloadReader.Factory payloadReaderFactory;
    private final SparseArray<TsPayloadReader> tsPayloadReaders; // Indexed by pid
    private final SparseBooleanArray trackIds;
    private final SparseBooleanArray trackPids;
    private final TsDurationReader durationReader;
    // Accessed only by the loading thread.
    private TsBinarySearchSeeker tsBinarySearchSeeker;
    private ExtractorOutput output;
    private int remainingPmts;
    private boolean tracksEnded;
    private boolean hasOutputSeekMap;
    private boolean pendingSeekToStart;
    @Nullable
    private TsPayloadReader id3Reader;
    private int bytesSinceLastSync;
    private int pcrPid;

    public MyTsExtractor() {
        this(/* defaultTsPayloadReaderFlags= */ 0);
    }

    /**
     * @param defaultTsPayloadReaderFlags A combination of {@link DefaultTsPayloadReaderFactory}
     *                                    {@code FLAG_*} values that control the behavior of the payload readers.
     */
    public MyTsExtractor(@TsPayloadReader.Flags int defaultTsPayloadReaderFlags) {
        this(MODE_SINGLE_PMT, defaultTsPayloadReaderFlags, DEFAULT_TIMESTAMP_SEARCH_BYTES);
    }

    /**
     * @param mode                        Mode for the extractor. One of {@link #MODE_MULTI_PMT}, {@link #MODE_SINGLE_PMT}
     *                                    and {@link #MODE_HLS}.
     * @param defaultTsPayloadReaderFlags A combination of {@link DefaultTsPayloadReaderFactory}
     *                                    {@code FLAG_*} values that control the behavior of the payload readers.
     * @param timestampSearchBytes        The number of bytes searched from a given position in the stream to
     *                                    find a PCR timestamp. If this value is too small, the duration might be unknown and seeking
     *                                    might not be supported for high bitrate progressive streams. Setting a large value for this
     *                                    field might be inefficient though because the extractor stores a buffer of {@code
     *                                    timestampSearchBytes} bytes when determining the duration or when performing a seek
     *                                    operation. The default value is {@link #DEFAULT_TIMESTAMP_SEARCH_BYTES}. If the number of
     *                                    bytes left in the stream from the current position is less than {@code
     *                                    timestampSearchBytes}, the search is performed on the bytes left.
     */
    public MyTsExtractor(
            @Mode int mode, @TsPayloadReader.Flags int defaultTsPayloadReaderFlags, int timestampSearchBytes) {
        this(
                mode,
                new TimestampAdjuster(0),
                new DefaultTsPayloadReaderFactory(defaultTsPayloadReaderFlags),
                timestampSearchBytes);
    }

    /**
     * @param mode                 Mode for the extractor. One of {@link #MODE_MULTI_PMT}, {@link #MODE_SINGLE_PMT}
     *                             and {@link #MODE_HLS}.
     * @param timestampAdjuster    A timestamp adjuster for offsetting and scaling sample timestamps.
     * @param payloadReaderFactory Factory for injecting a custom set of payload readers.
     */
    public MyTsExtractor(
            @Mode int mode,
            TimestampAdjuster timestampAdjuster,
            TsPayloadReader.Factory payloadReaderFactory) {
        this(mode, timestampAdjuster, payloadReaderFactory, DEFAULT_TIMESTAMP_SEARCH_BYTES);
    }

    /**
     * @param mode                 Mode for the extractor. One of {@link #MODE_MULTI_PMT}, {@link #MODE_SINGLE_PMT}
     *                             and {@link #MODE_HLS}.
     * @param timestampAdjuster    A timestamp adjuster for offsetting and scaling sample timestamps.
     * @param payloadReaderFactory Factory for injecting a custom set of payload readers.
     * @param timestampSearchBytes The number of bytes searched from a given position in the stream to
     *                             find a PCR timestamp. If this value is too small, the duration might be unknown and seeking
     *                             might not be supported for high bitrate progressive streams. Setting a large value for this
     *                             field might be inefficient though because the extractor stores a buffer of {@code
     *                             timestampSearchBytes} bytes when determining the duration or when performing a seek
     *                             operation. The default value is {@link #DEFAULT_TIMESTAMP_SEARCH_BYTES}. If the number of
     *                             bytes left in the stream from the current position is less than {@code
     *                             timestampSearchBytes}, the search is performed on the bytes left.
     */
    public MyTsExtractor(
            @Mode int mode,
            TimestampAdjuster timestampAdjuster,
            TsPayloadReader.Factory payloadReaderFactory,
            int timestampSearchBytes) {
        this.payloadReaderFactory = Assertions.checkNotNull(payloadReaderFactory);
        this.timestampSearchBytes = timestampSearchBytes;
        this.mode = mode;
        if (mode == MODE_SINGLE_PMT || mode == MODE_HLS) {
            timestampAdjusters = Collections.singletonList(timestampAdjuster);
        } else {
            timestampAdjusters = new ArrayList<>();
            timestampAdjusters.add(timestampAdjuster);
        }
        tsPacketBuffer = new ParsableByteArray(new byte[BUFFER_SIZE], 0);
        trackIds = new SparseBooleanArray();
        trackPids = new SparseBooleanArray();
        tsPayloadReaders = new SparseArray<>();
        continuityCounters = new SparseIntArray();
        durationReader = new TsDurationReader(timestampSearchBytes);
        output = ExtractorOutput.PLACEHOLDER;
        pcrPid = -1;
        resetPayloadReaders();
    }

    @Override
    public boolean sniff(ExtractorInput input) throws IOException {
        byte[] buffer = tsPacketBuffer.getData();

        input.peekFully(buffer, 0, TS_PACKET_SIZE * 2);
        for (int i = 0; i < TS_PACKET_SIZE * 2; i++) {
            if (buffer[i] == TS_SYNC_BYTE && buffer[i + 1] == 0x40) {
                if (i > 0) input.skipFully(i);
                break;
            }
        }
        input.peekFully(buffer, 0, TS_PACKET_SIZE * SNIFF_TS_PACKET_COUNT);
        for (int startPosCandidate = 0; startPosCandidate < TS_PACKET_SIZE; startPosCandidate++) {
            // Try to identify at least SNIFF_TS_PACKET_COUNT packets starting with TS_SYNC_BYTE.
            boolean isSyncBytePatternCorrect = true;
            for (int i = 0; i < SNIFF_TS_PACKET_COUNT; i++) {
                int ii = startPosCandidate + i * TS_PACKET_SIZE;
                if (buffer[ii] != TS_SYNC_BYTE) {
                    isSyncBytePatternCorrect = false;
                    break;
                }
            }
            if (isSyncBytePatternCorrect) {
                input.skipFully(startPosCandidate);
                return true;
            }
        }
        return false;
    }

    // Extractor implementation.

    @Override
    public void init(ExtractorOutput output) {
        this.output = output;
    }

    @Override
    public void seek(long position, long timeUs) {
        Assertions.checkState(mode != MODE_HLS);
        int timestampAdjustersCount = timestampAdjusters.size();
        for (int i = 0; i < timestampAdjustersCount; i++) {
            TimestampAdjuster timestampAdjuster = timestampAdjusters.get(i);
            // If the timestamp adjuster has not yet established a timestamp offset, we need to reset its
            // expected first sample timestamp to be the new seek position. Without this, the timestamp
            // adjuster would incorrectly establish its timestamp offset assuming that the first sample
            // after this seek corresponds to the start of the stream (or a previous seek position, if
            // there was one).
            boolean resetTimestampAdjuster = timestampAdjuster.getTimestampOffsetUs() == C.TIME_UNSET;
            if (!resetTimestampAdjuster) {
                long adjusterFirstSampleTimestampUs = timestampAdjuster.getFirstSampleTimestampUs();
                // Also reset the timestamp adjuster if its offset was calculated based on a non-zero
                // position in the stream (other than the position being seeked to), since in this case the
                // offset may not be accurate.
                resetTimestampAdjuster =
                        adjusterFirstSampleTimestampUs != C.TIME_UNSET
                                && adjusterFirstSampleTimestampUs != 0
                                && adjusterFirstSampleTimestampUs != timeUs;
            }
            if (resetTimestampAdjuster) {
                timestampAdjuster.reset(timeUs);
            }
        }
        if (timeUs != 0 && tsBinarySearchSeeker != null) {
            tsBinarySearchSeeker.setSeekTargetUs(timeUs);
        }
        tsPacketBuffer.reset(/* limit= */ 0);
        continuityCounters.clear();
        for (int i = 0; i < tsPayloadReaders.size(); i++) {
            tsPayloadReaders.valueAt(i).seek();
        }
        bytesSinceLastSync = 0;
    }

    @Override
    public void release() {
        // Do nothing
    }

    @Override
    @ReadResult
    public int read(ExtractorInput input, PositionHolder seekPosition) throws IOException {
        long inputLength = input.getLength();
        if (tracksEnded) {
            boolean canReadDuration = inputLength != C.LENGTH_UNSET && mode != MODE_HLS;
            if (canReadDuration && !durationReader.isDurationReadFinished()) {
                return durationReader.readDuration(input, seekPosition, pcrPid);
            }
            maybeOutputSeekMap(inputLength);

            if (pendingSeekToStart) {
                pendingSeekToStart = false;
                seek(/* position= */ 0, /* timeUs= */ 0);
                if (input.getPosition() != 0) {
                    seekPosition.position = 0;
                    return RESULT_SEEK;
                }
            }

            if (tsBinarySearchSeeker != null && tsBinarySearchSeeker.isSeeking()) {
                return tsBinarySearchSeeker.handlePendingSeek(input, seekPosition);
            }
        }

        if (!fillBufferWithAtLeastOnePacket(input)) {
            return RESULT_END_OF_INPUT;
        }

        int endOfPacket = findEndOfFirstTsPacketInBuffer();
        int limit = tsPacketBuffer.limit();
        if (endOfPacket > limit) {
            return RESULT_CONTINUE;
        }

        @TsPayloadReader.Flags int packetHeaderFlags = 0;

        // Note: See ISO/IEC 13818-1, section 2.4.3.2 for details of the header format.
        int tsPacketHeader = tsPacketBuffer.readInt();
        if ((tsPacketHeader & 0x800000) != 0) { // transport_error_indicator
            // There are uncorrectable errors in this packet.
            tsPacketBuffer.setPosition(endOfPacket);
            return RESULT_CONTINUE;
        }
        packetHeaderFlags |= (tsPacketHeader & 0x400000) != 0 ? FLAG_PAYLOAD_UNIT_START_INDICATOR : 0;
        // Ignoring transport_priority (tsPacketHeader & 0x200000)
        int pid = (tsPacketHeader & 0x1FFF00) >> 8;
        // Ignoring transport_scrambling_control (tsPacketHeader & 0xC0)
        boolean adaptationFieldExists = (tsPacketHeader & 0x20) != 0;
        boolean payloadExists = (tsPacketHeader & 0x10) != 0;

        TsPayloadReader payloadReader = payloadExists ? tsPayloadReaders.get(pid) : null;
        if (payloadReader == null) {
            tsPacketBuffer.setPosition(endOfPacket);
            return RESULT_CONTINUE;
        }

        // Discontinuity check.
        if (mode != MODE_HLS) {
            int continuityCounter = tsPacketHeader & 0xF;
            int previousCounter = continuityCounters.get(pid, continuityCounter - 1);
            continuityCounters.put(pid, continuityCounter);
            if (previousCounter == continuityCounter) {
                // Duplicate packet found.
                tsPacketBuffer.setPosition(endOfPacket);
                return RESULT_CONTINUE;
            } else if (continuityCounter != ((previousCounter + 1) & 0xF)) {
                // Discontinuity found.
                payloadReader.seek();
            }
        }

        // Skip the adaptation field.
        if (adaptationFieldExists) {
            int adaptationFieldLength = tsPacketBuffer.readUnsignedByte();
            int adaptationFieldFlags = tsPacketBuffer.readUnsignedByte();

            packetHeaderFlags |=
                    (adaptationFieldFlags & 0x40) != 0 // random_access_indicator.
                            ? TsPayloadReader.FLAG_RANDOM_ACCESS_INDICATOR
                            : 0;
            tsPacketBuffer.skipBytes(adaptationFieldLength - 1 /* flags */);
        }

        // Read the payload.
        boolean wereTracksEnded = tracksEnded;
        if (shouldConsumePacketPayload(pid)) {
            tsPacketBuffer.setLimit(endOfPacket);
            payloadReader.consume(tsPacketBuffer, packetHeaderFlags);
            tsPacketBuffer.setLimit(limit);
        }
        if (mode != MODE_HLS && !wereTracksEnded && tracksEnded && inputLength != C.LENGTH_UNSET) {
            // We have read all tracks from all PMTs in this non-live stream. Now seek to the beginning
            // and read again to make sure we output all media, including any contained in packets prior
            // to those containing the track information.
            pendingSeekToStart = true;
        }

        tsPacketBuffer.setPosition(endOfPacket);
        return RESULT_CONTINUE;
    }

    private void maybeOutputSeekMap(long inputLength) {
        if (!hasOutputSeekMap) {
            hasOutputSeekMap = true;
            if (durationReader.getDurationUs() != C.TIME_UNSET) {
                tsBinarySearchSeeker =
                        new TsBinarySearchSeeker(
                                durationReader.getPcrTimestampAdjuster(),
                                durationReader.getDurationUs(),
                                inputLength,
                                pcrPid,
                                timestampSearchBytes);
                output.seekMap(tsBinarySearchSeeker.getSeekMap());
            } else {
                output.seekMap(new SeekMap.Unseekable(durationReader.getDurationUs()));
            }
        }
    }

    // Internals.

    private boolean fillBufferWithAtLeastOnePacket(ExtractorInput input) throws IOException {
        byte[] data = tsPacketBuffer.getData();
        // Shift bytes to the start of the buffer if there isn't enough space left at the end.
        if (BUFFER_SIZE - tsPacketBuffer.getPosition() < TS_PACKET_SIZE) {
            int bytesLeft = tsPacketBuffer.bytesLeft();
            if (bytesLeft > 0) {
                System.arraycopy(data, tsPacketBuffer.getPosition(), data, 0, bytesLeft);
            }
            tsPacketBuffer.reset(data, bytesLeft);
        }
        // Read more bytes until we have at least one packet.
        while (tsPacketBuffer.bytesLeft() < TS_PACKET_SIZE) {
            int limit = tsPacketBuffer.limit();
            int read = input.read(data, limit, BUFFER_SIZE - limit);
            if (read == C.RESULT_END_OF_INPUT) {
                return false;
            }
            tsPacketBuffer.setLimit(limit + read);
        }
        return true;
    }

    /**
     * Returns the position of the end of the first TS packet (exclusive) in the packet buffer.
     *
     * <p>This may be a position beyond the buffer limit if the packet has not been read fully into
     * the buffer, or if no packet could be found within the buffer.
     */
    private int findEndOfFirstTsPacketInBuffer() throws ParserException {
        int searchStart = tsPacketBuffer.getPosition();
        int limit = tsPacketBuffer.limit();
        int syncBytePosition =
                TsUtil.findSyncBytePosition(tsPacketBuffer.getData(), searchStart, limit);
        // Discard all bytes before the sync byte.
        // If sync byte is not found, this means discard the whole buffer.
        tsPacketBuffer.setPosition(syncBytePosition);
        int endOfPacket = syncBytePosition + TS_PACKET_SIZE;
        if (endOfPacket > limit) {
            bytesSinceLastSync += syncBytePosition - searchStart;
            if (mode == MODE_HLS && bytesSinceLastSync > TS_PACKET_SIZE * 2) {
                throw ParserException.createForMalformedContainer(
                        "Cannot find sync byte. Most likely not a Transport Stream.", /* cause= */ null);
            }
        } else {
            // We have found a packet within the buffer.
            bytesSinceLastSync = 0;
        }
        return endOfPacket;
    }

    private boolean shouldConsumePacketPayload(int packetPid) {
        return mode == MODE_HLS
                || tracksEnded
                || !trackPids.get(packetPid, /* valueIfKeyNotFound= */ false); // It's a PSI packet
    }

    private void resetPayloadReaders() {
        trackIds.clear();
        tsPayloadReaders.clear();
        SparseArray<TsPayloadReader> initialPayloadReaders =
                payloadReaderFactory.createInitialPayloadReaders();
        int initialPayloadReadersSize = initialPayloadReaders.size();
        for (int i = 0; i < initialPayloadReadersSize; i++) {
            tsPayloadReaders.put(initialPayloadReaders.keyAt(i), initialPayloadReaders.valueAt(i));
        }
        tsPayloadReaders.put(TS_PAT_PID, new SectionReader(new PatReader()));
        id3Reader = null;
    }

    /**
     * Modes for the extractor. One of {@link #MODE_MULTI_PMT}, {@link #MODE_SINGLE_PMT} or {@link
     * #MODE_HLS}.
     */
    @Documented
    @Retention(RetentionPolicy.SOURCE)
    @IntDef({MODE_MULTI_PMT, MODE_SINGLE_PMT, MODE_HLS})
    public @interface Mode {
    }

    /**
     * Parses Program Association Table data.
     */
    private class PatReader implements SectionPayloadReader {

        private final ParsableBitArray patScratch;

        public PatReader() {
            patScratch = new ParsableBitArray(new byte[4]);
        }

        @Override
        public void init(
                TimestampAdjuster timestampAdjuster,
                ExtractorOutput extractorOutput,
                TsPayloadReader.TrackIdGenerator idGenerator) {
            // Do nothing.
        }

        @Override
        public void consume(ParsableByteArray sectionData) {
            int tableId = sectionData.readUnsignedByte();
            if (tableId != 0x00 /* program_association_section */) {
                // See ISO/IEC 13818-1, section 2.4.4.4 for more information on table id assignment.
                return;
            }
            // section_syntax_indicator(1), '0'(1), reserved(2), section_length(4)
            int secondHeaderByte = sectionData.readUnsignedByte();
            if ((secondHeaderByte & 0x80) == 0) {
                // section_syntax_indicator must be 1. See ISO/IEC 13818-1, section 2.4.4.5.
                return;
            }
            // section_length(8), transport_stream_id (16), reserved (2), version_number (5),
            // current_next_indicator (1), section_number (8), last_section_number (8)
            sectionData.skipBytes(6);

            int programCount = sectionData.bytesLeft() / 4;
            for (int i = 0; i < programCount; i++) {
                sectionData.readBytes(patScratch, 4);
                int programNumber = patScratch.readBits(16);
                patScratch.skipBits(3); // reserved (3)
                if (programNumber == 0) {
                    patScratch.skipBits(13); // network_PID (13)
                } else {
                    int pid = patScratch.readBits(13);
                    if (tsPayloadReaders.get(pid) == null) {
                        tsPayloadReaders.put(pid, new SectionReader(new PmtReader(pid)));
                        remainingPmts++;
                    }
                }
            }
            if (mode != MODE_HLS) {
                tsPayloadReaders.remove(TS_PAT_PID);
            }
        }
    }

    /**
     * Parses Program Map Table.
     */
    private class PmtReader implements SectionPayloadReader {

        private static final int TS_PMT_DESC_REGISTRATION = 0x05;
        private static final int TS_PMT_DESC_ISO639_LANG = 0x0A;
        private static final int TS_PMT_DESC_AC3 = 0x6A;
        private static final int TS_PMT_DESC_AIT = 0x6F;
        private static final int TS_PMT_DESC_EAC3 = 0x7A;
        private static final int TS_PMT_DESC_DTS = 0x7B;
        private static final int TS_PMT_DESC_DVB_EXT = 0x7F;
        private static final int TS_PMT_DESC_DVBSUBS = 0x59;

        private static final int TS_PMT_DESC_DVB_EXT_AC4 = 0x15;

        private final ParsableBitArray pmtScratch;
        private final SparseArray<TsPayloadReader> trackIdToReaderScratch;
        private final SparseIntArray trackIdToPidScratch;
        private final int pid;

        public PmtReader(int pid) {
            pmtScratch = new ParsableBitArray(new byte[5]);
            trackIdToReaderScratch = new SparseArray<>();
            trackIdToPidScratch = new SparseIntArray();
            this.pid = pid;
        }

        @Override
        public void init(
                TimestampAdjuster timestampAdjuster,
                ExtractorOutput extractorOutput,
                TsPayloadReader.TrackIdGenerator idGenerator) {
            // Do nothing.
        }

        @Override
        public void consume(ParsableByteArray sectionData) {
            int tableId = sectionData.readUnsignedByte();
            if (tableId != 0x02 /* TS_program_map_section */) {
                // See ISO/IEC 13818-1, section 2.4.4.4 for more information on table id assignment.
                return;
            }
            // TimestampAdjuster assignment.
            TimestampAdjuster timestampAdjuster;
            if (mode == MODE_SINGLE_PMT || mode == MODE_HLS || remainingPmts == 1) {
                timestampAdjuster = timestampAdjusters.get(0);
            } else {
                timestampAdjuster =
                        new TimestampAdjuster(timestampAdjusters.get(0).getFirstSampleTimestampUs());
                timestampAdjusters.add(timestampAdjuster);
            }

            // section_syntax_indicator(1), '0'(1), reserved(2), section_length(4)
            int secondHeaderByte = sectionData.readUnsignedByte();
            if ((secondHeaderByte & 0x80) == 0) {
                // section_syntax_indicator must be 1. See ISO/IEC 13818-1, section 2.4.4.9.
                return;
            }
            // section_length(8)
            sectionData.skipBytes(1);
            int programNumber = sectionData.readUnsignedShort();

            // Skip 3 bytes (24 bits), including:
            // reserved (2), version_number (5), current_next_indicator (1), section_number (8),
            // last_section_number (8)
            sectionData.skipBytes(3);

            sectionData.readBytes(pmtScratch, 2);
            // reserved (3), PCR_PID (13)
            pmtScratch.skipBits(3);
            pcrPid = pmtScratch.readBits(13);

            // Read program_info_length.
            sectionData.readBytes(pmtScratch, 2);
            pmtScratch.skipBits(4);
            int programInfoLength = pmtScratch.readBits(12);

            // Skip the descriptors.
            sectionData.skipBytes(programInfoLength);

            if (mode == MODE_HLS && id3Reader == null) {
                // Setup an ID3 track regardless of whether there's a corresponding entry, in case one
                // appears intermittently during playback. See [Internal: b/20261500].
                TsPayloadReader.EsInfo id3EsInfo = new TsPayloadReader.EsInfo(TS_STREAM_TYPE_ID3, null, AUDIO_TYPE_UNDEFINED,null, Util.EMPTY_BYTE_ARRAY);
                id3Reader = payloadReaderFactory.createPayloadReader(TS_STREAM_TYPE_ID3, id3EsInfo);
                if (id3Reader != null) {
                    id3Reader.init(
                            timestampAdjuster,
                            output,
                            new TsPayloadReader.TrackIdGenerator(programNumber, TS_STREAM_TYPE_ID3, MAX_PID_PLUS_ONE));
                }
            }

            trackIdToReaderScratch.clear();
            trackIdToPidScratch.clear();
            int remainingEntriesLength = sectionData.bytesLeft();
            while (remainingEntriesLength > 0) {
                sectionData.readBytes(pmtScratch, 5);
                int streamType = pmtScratch.readBits(8);
                pmtScratch.skipBits(3); // reserved
                int elementaryPid = pmtScratch.readBits(13);
                pmtScratch.skipBits(4); // reserved
                int esInfoLength = pmtScratch.readBits(12); // ES_info_length.
                TsPayloadReader.EsInfo esInfo = readEsInfo(sectionData, esInfoLength);
                if (streamType == 0x06 || streamType == 0x05) {
                    streamType = esInfo.streamType;
                }
                remainingEntriesLength -= esInfoLength + 5;

                int trackId = mode == MODE_HLS ? streamType : elementaryPid;
                if (trackIds.get(trackId)) {
                    continue;
                }

                @Nullable
                TsPayloadReader reader =
                        mode == MODE_HLS && streamType == TS_STREAM_TYPE_ID3
                                ? id3Reader
                                : payloadReaderFactory.createPayloadReader(streamType, esInfo);
                if (mode != MODE_HLS
                        || elementaryPid < trackIdToPidScratch.get(trackId, MAX_PID_PLUS_ONE)) {
                    trackIdToPidScratch.put(trackId, elementaryPid);
                    trackIdToReaderScratch.put(trackId, reader);
                }
            }

            int trackIdCount = trackIdToPidScratch.size();
            for (int i = 0; i < trackIdCount; i++) {
                int trackId = trackIdToPidScratch.keyAt(i);
                int trackPid = trackIdToPidScratch.valueAt(i);
                trackIds.put(trackId, true);
                trackPids.put(trackPid, true);
                @Nullable TsPayloadReader reader = trackIdToReaderScratch.valueAt(i);
                if (reader != null) {
                    if (reader != id3Reader) {
                        reader.init(
                                timestampAdjuster,
                                output,
                                new TsPayloadReader.TrackIdGenerator(programNumber, trackId, MAX_PID_PLUS_ONE));
                    }
                    tsPayloadReaders.put(trackPid, reader);
                }
            }

            if (mode == MODE_HLS) {
                if (!tracksEnded) {
                    output.endTracks();
                    remainingPmts = 0;
                    tracksEnded = true;
                }
            } else {
                tsPayloadReaders.remove(pid);
                remainingPmts = mode == MODE_SINGLE_PMT ? 0 : remainingPmts - 1;
                if (remainingPmts == 0) {
                    output.endTracks();
                    tracksEnded = true;
                }
            }
        }

        /**
         * Returns the stream info read from the available descriptors. Sets {@code data}'s position to
         * the end of the descriptors.
         *
         * @param data   A buffer with its position set to the start of the first descriptor.
         * @param length The length of descriptors to read from the current position in {@code data}.
         * @return The stream info read from the available descriptors.
         */
        private TsPayloadReader.EsInfo readEsInfo(ParsableByteArray data, int length) {
            int descriptorsStartPosition = data.getPosition();
            int descriptorsEndPosition = descriptorsStartPosition + length;
            int streamType = -1;
            @TsPayloadReader.EsInfo.AudioType int audioType = AUDIO_TYPE_UNDEFINED;
            String language = null;
            List<TsPayloadReader.DvbSubtitleInfo> dvbSubtitleInfos = null;
            while (data.getPosition() < descriptorsEndPosition) {
                int descriptorTag = data.readUnsignedByte();
                int descriptorLength = data.readUnsignedByte();
                int positionOfNextDescriptor = data.getPosition() + descriptorLength;
                if (positionOfNextDescriptor > descriptorsEndPosition) {
                    // Descriptor claims to extend past the end position. Skip it.
                    break;
                }
                if (descriptorTag == TS_PMT_DESC_REGISTRATION) { // registration_descriptor
                    long formatIdentifier = data.readUnsignedInt();
                    if (formatIdentifier == AC3_FORMAT_IDENTIFIER) {
                        streamType = TS_STREAM_TYPE_AC3;
                    } else if (formatIdentifier == E_AC3_FORMAT_IDENTIFIER) {
                        streamType = TS_STREAM_TYPE_E_AC3;
                    } else if (formatIdentifier == AC4_FORMAT_IDENTIFIER) {
                        streamType = TS_STREAM_TYPE_AC4;
                    } else if (formatIdentifier == HEVC_FORMAT_IDENTIFIER) {
                        streamType = TS_STREAM_TYPE_H265;
                    }
                } else if (descriptorTag == TS_PMT_DESC_AC3) { // AC-3_descriptor in DVB (ETSI EN 300 468)
                    streamType = TS_STREAM_TYPE_AC3;
                } else if (descriptorTag == TS_PMT_DESC_EAC3) { // enhanced_AC-3_descriptor
                    streamType = TS_STREAM_TYPE_E_AC3;
                } else if (descriptorTag == TS_PMT_DESC_DVB_EXT) {
                    // Extension descriptor in DVB (ETSI EN 300 468).
                    int descriptorTagExt = data.readUnsignedByte();
                    if (descriptorTagExt == TS_PMT_DESC_DVB_EXT_AC4) {
                        // AC-4_descriptor in DVB (ETSI EN 300 468).
                        streamType = TS_STREAM_TYPE_AC4;
                    }
                } else if (descriptorTag == TS_PMT_DESC_DTS) { // DTS_descriptor
                    streamType = TS_STREAM_TYPE_DTS;
                } else if (descriptorTag == TS_PMT_DESC_ISO639_LANG) {
                    language = data.readString(3).trim();
                    audioType = data.readUnsignedByte();
                } else if (descriptorTag == TS_PMT_DESC_DVBSUBS) {
                    streamType = TS_STREAM_TYPE_DVBSUBS;
                    dvbSubtitleInfos = new ArrayList<>();
                    while (data.getPosition() < positionOfNextDescriptor) {
                        String dvbLanguage = data.readString(3).trim();
                        int dvbSubtitlingType = data.readUnsignedByte();
                        byte[] initializationData = new byte[4];
                        data.readBytes(initializationData, 0, 4);
                        dvbSubtitleInfos.add(
                                new TsPayloadReader.DvbSubtitleInfo(dvbLanguage, dvbSubtitlingType, initializationData));
                    }
                } else if (descriptorTag == TS_PMT_DESC_AIT) {
                    streamType = TS_STREAM_TYPE_AIT;
                }
                // Skip unused bytes of current descriptor.
                data.skipBytes(positionOfNextDescriptor - data.getPosition());
            }
            data.setPosition(descriptorsEndPosition);
            return new TsPayloadReader.EsInfo(
                    streamType,
                    language,
                    audioType,
                    dvbSubtitleInfos,
                    Arrays.copyOfRange(data.getData(), descriptorsStartPosition, descriptorsEndPosition));
        }

    }

}